use rand::Rng;
use aes::Aes256;
use aes::cipher::{KeyIvInit, StreamCipher};
use std::convert::TryInto;
use hex::{encode, decode};

type Aes256Ctr = ctr::Ctr128BE<Aes256>;


fn encrypt_hexa(keys: Vec<u8>) -> String{
    let hex_keys = encode(&keys);
    println!("-------");
    println!("Encrypted (HeX): {}", hex_keys);
    println!("-------");

    hex_keys
}

fn decrypt_hexa(keys: String)-> Vec<u8>{
    let dec_hexa = decode(keys).expect("Failed to decode Hex");
    let decrypt_vec_key: Vec<u8> = dec_hexa.iter().cloned().collect();
    println!("------");
    println!("DeC Keys: {:?}", decrypt_vec_key);
    println!("------");

    decrypt_vec_key
}

fn main() {
    // Generate random keys and nonce
    let aes_key: [u8; 32] = rand::thread_rng().gen();
    let chacha_key: [u8; 32] = rand::thread_rng().gen();
    let chacha_nonce: [u8; 12] = rand::thread_rng().gen();

    println!("AES Key without HEx: {:?}", aes_key);
    println!("AES Key: {:x?}", aes_key);
    println!("ChaCha Key: {:?}", chacha_key);
    println!("ChaCha Nonce: {:?}", chacha_nonce);

    // Combine the keys into a single key
    let mut combined_key = Vec::new();
    combined_key.extend_from_slice(&aes_key);
    combined_key.extend_from_slice(&chacha_key);
    combined_key.extend_from_slice(&chacha_nonce);

    // Encrypt the combined key to create the main key
    let main_key = encrypt_combined_key(&combined_key, &aes_key);

    println!();
    println!("Main Key: {:?}", main_key);

    // Encrypt the main key in Hexa
    let main_key_hex = encrypt_hexa(main_key.clone());


    // Decrypt from hex
    let decrypted_keys = decrypt_hexa(main_key_hex);

    
    // Decrypt from the main key and extract the original one ..!
    let decrypted_keys = decrypt_combined_key(&decrypted_keys, &aes_key).unwrap();
    let (decrypted_aes_key, decrypted_chacha_key, decrypted_chacha_nonce) = split_keys(&decrypted_keys);

    println!("Decrypted AES Key: {:x?}", decrypted_aes_key);
    println!("Decrypted ChaCha Key: {:x?}", decrypted_chacha_key);
    println!("Decrypted ChaCha Nonce: {:x?}", decrypted_chacha_nonce);
}

// Encrypt the combined key
fn encrypt_combined_key(combined_key: &[u8], aes_key: &[u8; 32]) -> Vec<u8> {
    let aes_iv: [u8; 16] = rand::thread_rng().gen(); // Generate a random IV
    let mut cipher = Aes256Ctr::new(aes_key.into(), &aes_iv.into());
    let mut encrypted_key = combined_key.to_vec();
    cipher.apply_keystream(&mut encrypted_key);
    let mut result = Vec::new();
    result.extend_from_slice(&aes_iv);
    result.extend_from_slice(&encrypted_key);
    result
}

// Decrypt the main key
fn decrypt_combined_key(encrypted_key: &[u8], aes_key: &[u8; 32]) -> Result<Vec<u8>, &'static str> {
    if encrypted_key.len() < 16 {
        return Err("Invalid encrypted key length");
    }
    let aes_iv: [u8; 16] = encrypted_key[0..16].try_into().unwrap();
    let mut cipher = Aes256Ctr::new(aes_key.into(), &aes_iv.into());
    let mut decrypted_key = encrypted_key[16..].to_vec();
    cipher.apply_keystream(&mut decrypted_key);
    Ok(decrypted_key)
}

// Split the decrypted keys into the original components
fn split_keys(combined_key: &[u8]) -> ([u8; 32], [u8; 32], [u8; 12]) {
    let aes_key: [u8; 32] = combined_key[0..32].try_into().unwrap();
    let chacha_key: [u8; 32] = combined_key[32..64].try_into().unwrap();
    let chacha_nonce: [u8; 12] = combined_key[64..76].try_into().unwrap();
    (aes_key, chacha_key, chacha_nonce)
}




